import traffic_generator
from itertools import product, combinations
import numpy as np

from numpy.random import seed
from numpy.random import rand

'''
Splits the set into two equal or almost-equal parts, and then sends from every node
in one set to all nodes in the second set.
'''
class BipartiteTrafficGenerator(traffic_generator.TrafficGenerator):
	def __init__(self, nblocks):
		traffic_generator.TrafficGenerator.__init__(self, nblocks)
		return

	def get_name(self):
		return "bipartite"

	def _generate_number_of_snapshots(self, set1_block_ids, nsnapshots):
		seed(1)
		all_tm_snapshots = []
		set1_size = len(set1_block_ids)
		set2_size = self.nblocks - set1_size
		for index in range(nsnapshots):
			random_values = rand(set1_size)
			direction = np.random.choice([-1, 1], size=(10,), p=[1./2, 1./2])
			sum_values = sum(random_values)
			outgoing_traffic = [float(x)/sum_values for x in random_values]
			tm = np.zeros((self.nblocks, self.nblocks,))
			for block_index in set1_block_ids:
				for potential_target_block in range(self.nblocks):
					if potential_target_block not in set1_block_ids:
						if direction[block_index] > 0:
							tm[block_index][potential_target_block] = outgoing_traffic[block_index] / set2_size
						else:
							tm[potential_target_block][block_index] = outgoing_traffic[block_index] / set2_size
			all_tm_snapshots.append(tm)
		return all_tm_snapshots

	## returns the probability communication matrix such that all entries in the matrix
	## sums to 1
	def generate_probability_traffic_matrices(self, nsnapshots):
		set1_size = self.nblocks / 2
		set2_size = self.nblocks - set1_size
		# find all the combinations for 
		block_index_sets = range(self.nblocks)
		set1_combinations_of_blocks = combinations(block_index_sets, set1_size)
		return self._generate_number_of_snapshots(set(range(set1_size)), nsnapshots)